from time import sleep
from math import pi
import numpy as np
import rospy
from unitree_legged_msgs.msg import LowCmd, LowState, MotorCmd, MotorState

from low_state_subscriber import LowStateSubscriber
from low_cmd_publisher import LowCmdPublisher


class Controller():
    def __init__(self, rname):
        self.robot_name = rname
        self.state_subscriber = LowStateSubscriber(self.robot_name)
        self.cmd_publisher = LowCmdPublisher(self.robot_name)
        self.lowCmd = LowCmd()

    def paramInit(self):
        for i in range(4):
            self.lowCmd.motorCmd[i*3+0].mode = 0x0A
            self.lowCmd.motorCmd[i*3+0].Kp = 70     # desired position stiffness位置刚度 (unit: N.m/rad )
            self.lowCmd.motorCmd[i*3+0].dq = 0      # desired velocity (unit: radian/second)
            self.lowCmd.motorCmd[i*3+0].Kd = 3      # desired velocity stiffness速度刚度 (unit: N.m/(rad/s) )
            self.lowCmd.motorCmd[i*3+0].tau = 0     # desired output torque扭矩 (unit: N.m)
            self.lowCmd.motorCmd[i*3+1].mode = 0x0A
            self.lowCmd.motorCmd[i*3+1].Kp = 180
            self.lowCmd.motorCmd[i*3+1].dq = 0
            self.lowCmd.motorCmd[i*3+1].Kd = 8
            self.lowCmd.motorCmd[i*3+1].tau = 0
            self.lowCmd.motorCmd[i*3+2].mode = 0x0A
            self.lowCmd.motorCmd[i*3+2].Kp = 300
            self.lowCmd.motorCmd[i*3+2].dq = 0
            self.lowCmd.motorCmd[i*3+2].Kd = 15
            self.lowCmd.motorCmd[i*3+2].tau = 0
        for i in range(12):
            self.lowCmd.motorCmd[i].q = self.state_subscriber.lowState.motorState[i].q   # current angle当前角度 (unit: radian)

    def motion_init(self):
        self.paramInit()
        self.down()
        self.stand()

    def takePosition(self, targetPos, duration, motor_index=np.arange(12)):
        lastPos = list(np.arange(12))
        for i in range(12):
            lastPos[i] = self.state_subscriber.lowState.motorState[i].q
        for i in range(duration):
            if rospy.is_shutdown(): break
            percent = i/duration
            for j in motor_index:
                self.lowCmd.motorCmd[j].q = lastPos[j]*(1-percent) + targetPos[j]*percent
            self.cmd_publisher.sendServoCmd(self.lowCmd)




    def down(self):
        pos = [0.0, pi/3, -2.6,  -0.0, pi/3, -2.6,
               0.0, pi/3, -2.6,  -0.0, pi/3, -2.6]
        self.movePosition(pos, 300)

    def stand(self):
        pos = [0.0, pi/4, -pi/2,  -0.0, pi/4, -pi/2, 
               0.0, pi/4, -pi/2,  -0.0, pi/4, -pi/2]
        self.movePosition(pos, 300)